import baostock as bs
import pandas as pd


def write_one_stock(symbol: str, start_date: str, end_date: str):
    fields = (
        "date,code,open,high,low,close,volume,amount,"
        "adjustflag,turn,tradestatus,pctChg,peTTM,"
        "pbMRQ,psTTM,pcfNcfTTM,isST"
    )

    #### 登陆系统 ####
    lg = bs.login()
    if lg.error_code != "0":
        print("login respond error_code:" + lg.error_code)
        print("login respond  error_msg:" + lg.error_msg)

    #### 获取沪深A股历史K线数据 ####
    # 详细指标参数，参见“历史行情指标参数”章节；“分钟线”参数与“日线”参数不同。“分钟线”不包含指数。
    # 分钟线指标：date,time,code,open,high,low,close,volume,amount,adjustflag
    # 周月线指标：date,code,open,high,low,close,volume,amount,adjustflag,turn,pctChg
    rs = bs.query_history_k_data_plus(
        symbol,
        fields,
        start_date=start_date,
        end_date=end_date,
        frequency="d",
        adjustflag="2",
    )
    if rs.error_code != "0":
        print("login respond error_code:" + rs.error_code)
        print("login respond  error_msg:" + rs.error_msg)

    data_list = []
    while (rs.error_code == "0") & rs.next():
        data_list.append(rs.get_row_data())
    result = pd.DataFrame(data_list, columns=rs.fields)

    #### 结果集输出到csv文件 ####
    # result.to_csv("D:\\history_A_stock_k_data.csv", index=False)
    # print(result)

    #### 登出系统 ####
    bs.logout()


if __name__ == "__main__":
    write_one_stock("sh.600000", "2015-01-01", "2024-01-01")
